"""多智能体版本：
- ReactIAgent: 负责逐步 [Think]/[Action]/[Observation]
- SummaryAgent: 整合过程，输出 [Summary]/[FinalAnswer]
- OutputProcessor: 处理不同输出类型需求

运行：
    python react_agent.py
"""
from __future__ import annotations

import asyncio
import json
import re
import textwrap
import os
from datetime import datetime
from typing import Annotated, Optional, Literal
from typing_extensions import TypedDict
import aiohttp
from pathlib import Path

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver

# ---------- 输出处理器模块 ----------
"""专业输出处理器"""
class ProfessionalOutputProcessor:
    
    def __init__(self, output_dir: Path):
        self.output_dir = output_dir
        self.output_dir.mkdir(exist_ok=True)
        self.references = []
    
    def process_output(self, content: str, output_type: str, query: str) -> Tuple[str, Optional[str]]:
        """处理输出内容并生成专业文件"""
        final_content = self.extract_final_content(content)
        file_path = None
        
        if output_type == "markdown":
            file_path = self._create_professional_markdown(final_content, query)
        
        return final_content, file_path
    
    def extract_final_content(self, content: str) -> str:
        """提取最终答案内容"""
        # 先尝试匹配[FinalAnswer]标签
        final_match = re.search(r'\[FinalAnswer\](.*?)(?=\[|$)', content, re.DOTALL)
        if final_match:
            return final_match.group(1).strip()
        
        # 如果没有[FinalAnswer]标签，尝试匹配[任务完成]标签
        task_done_match = re.search(r'\[任务完成\](.*?)(?=\[|$)', content, re.DOTALL)
        if task_done_match:
            return task_done_match.group(1).strip()
        
        # 如果都没有匹配到，返回原始内容
        return content
    
    def _create_professional_markdown(self, content: str, query: str) -> str:
        """创建专业Markdown报告"""
        # 添加报告标题
        md_content = f"# {query}\n\n"
        md_content += f"**生成日期**: {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n"
        md_content += "---\n\n"
        md_content += content
        
        # 保存文件
        filename = f"professional_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
        filepath = self.output_dir / filename
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(md_content)
        
        return str(filepath)


# ---------- 输出提示词模块（合并自output_prompts.py） ----------
def build_summary_prompt(output_type: str, task: str, files: list) -> str:
    """根据输出类型构建提示词"""
    base_prompt = ""
    
    if output_type == "markdown":
        base_prompt = build_markdown_prompt(task, files)
    else:
        base_prompt = (
            "请整合整个对话过程并生成最终答案。\n"
            "输出格式要求：\n"
            "[FinalAnswer] 给用户的最终答案（直接、简洁、完整）\n\n"
            "重要提示：\n"
            "• 不要重复已输出的任何内容\n"
            "• 确保最终答案包含所有必要信息，避免冗余细节\n"
        )
    
    return base_prompt

def build_markdown_prompt(task: str, files: list) -> str:
    """构建专业Markdown提示词"""
    return f"""
    ## 专业报告生成规范
    
    ### 核心要求
    - 语言：中文（专有名词保留原文）
    - 信息来源：严格基于提供内容（禁止编造）
    - 标注规范：[[编号]](链接) 格式标注来源
    
    ### 报告结构
    1. 结构规划：金字塔原理（结论→细节）
    2. 内容提取：关键结论+背景数据+多视角分析
    3. 争议处理：客观呈现不同观点
    4. 数据深度：多维度交叉分析
    
    ### 格式要求
    - 标题层级： # → ## → ###
    - 表格：Markdown表格（带总结）
    - 公式：LaTeX格式
    - 图表：mermaid语法
    - 禁止：<a>标签
    
    ### 内容要求
    - 最小长度：5万字
    - 包含实用信息：工具推荐/常见误区
    
    ## 任务内容
    {task}
    
    **生成详细专业报告**
    """

def format_file_list(files: list) -> str:
    """格式化文件列表"""
    if not files:
        return "无"
        
    return "\n".join(
        f"- [{f.get('title', f.get('description', 'Untitled'))}]({f.get('link', '')})"
        for f in files
    )


# ---------- 配置 ----------
def load_cfg(path: str = "api_config_local.json"):
    with open(path, encoding="utf-8") as f:
        return json.load(f)


cfg = load_cfg()
llm = ChatOpenAI(
    model=cfg["model"],
    temperature=cfg.get("temperature", 0),
    api_key=cfg.get("api_key", "EMPTY"),
    base_url=cfg["base_url"],
    streaming=True,
)

MCP_CLIENT_URL = "http://localhost:7125"
OUTPUT_DIR = Path("outputs")

# 创建输出处理器实例
output_processor = ProfessionalOutputProcessor(OUTPUT_DIR)

# ---------- 辅助：从文本中查找 JSON ----------
def find_json_in_text(text: str) -> Optional[str]:
    """尝试从文本中提取第一个 JSON 对象字符串"""
    if not text:
        return None
    m = re.search(r'```json\s*(\{.*?\})\s*```', text, re.S)
    if m:
        return m.group(1)
    m = re.search(r'```(?:.*?)?(\{.*?\})(?:.*?)?```', text, re.S)
    if m:
        return m.group(1)
    start = text.find("{")
    if start == -1:
        return None
    stack, in_str, esc = 0, False, False
    for i in range(start, len(text)):
        ch = text[i]
        if ch == '"' and not esc:
            in_str = not in_str
        if in_str and ch == "\\" and not esc:
            esc = True
            continue
        if esc:
            esc = False
            continue
        if not in_str:
            if ch == "{":
                stack += 1
            elif ch == "}":
                stack -= 1
                if stack == 0:
                    return text[start:i + 1]
    return None


# ---------- 状态 ----------
class State(TypedDict):
    messages: Annotated[list, add_messages]
    output_type: Literal["normal", "markdown"]
    output_file: Optional[str]


# ---------- 连接状态检查 ----------
async def check_mcp_connection() -> bool:
    """检查MCP客户端网关的连接状态"""
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(f"{MCP_CLIENT_URL}/ping", timeout=10) as resp:
                return resp.status == 200
    except:
        return False

# ---------- 工具获取 ----------
class MCPConnectionManager:
    """MCP连接管理器，优化性能和缓存策略"""
    _connection_cache = None
    _tools_cache = None
    _session_pool = None
    _cache_timeout = 300  # 5分钟缓存
    _last_check_time = 0
    _connection_pool_timeout = 600  # 连接池超时时间（秒）
    
    @classmethod
    async def _get_session(cls):
        """获取或创建连接会话，支持连接池"""
        import time
        current_time = time.time()
        
        # 如果连接池不存在或已过期，创建新的会话
        if (cls._session_pool is None or 
            current_time - cls._last_check_time > cls._connection_pool_timeout):
            # 确保之前的会话已关闭
            if cls._session_pool:
                await cls._session_pool.close()
            cls._session_pool = aiohttp.ClientSession()
            cls._last_check_time = current_time
        
        return cls._session_pool
    
    @classmethod
    async def check_connection(cls) -> bool:
        """检查连接状态，带智能缓存机制"""
        import time
        current_time = time.time()
        
        # 如果缓存有效且未过期，直接返回缓存结果
        if cls._connection_cache is not None and current_time - cls._last_check_time < cls._cache_timeout:
            return cls._connection_cache
        
        # 实际检查连接，使用连接池
        try:
            session = await cls._get_session()
            async with session.get(f"{MCP_CLIENT_URL}/ping", timeout=10) as resp:
                cls._connection_cache = resp.status == 200
                cls._last_check_time = current_time
                return cls._connection_cache
        except:
            cls._connection_cache = False
            cls._last_check_time = current_time
            return False
    
    @classmethod
    async def cleanup(cls):
        """清理连接池资源"""
        if cls._session_pool:
            await cls._session_pool.close()
            cls._session_pool = None
            cls._connection_cache = None
            cls._tools_cache = None

async def fetch_tools(max_retries: int = 3) -> dict:
    """从MCP客户端网关获取工具列表，优化性能和缓存策略"""
    import time
    current_time = time.time()
    
    # 检查缓存是否有效
    if (MCPConnectionManager._tools_cache is not None and 
        current_time - MCPConnectionManager._last_check_time < MCPConnectionManager._cache_timeout):
        return MCPConnectionManager._tools_cache
    
    for attempt in range(max_retries):
        try:
            # 使用连接管理器检查连接状态
            if not await MCPConnectionManager.check_connection():
                if attempt < max_retries - 1:
                    await asyncio.sleep(2 ** attempt)  # 指数退避
                    continue
                result = {"error": "MCP客户端网关连接失败，请检查服务是否启动"}
                MCPConnectionManager._tools_cache = result
                return result
            
            # 使用连接池获取会话
            session = await MCPConnectionManager._get_session()
            async with session.get(f"{MCP_CLIENT_URL}/tools", timeout=30) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    tools_list = data.get("tools", [])
                    
                    if not tools_list:
                        result = {"error": "MCP服务器未提供任何工具"}
                        MCPConnectionManager._tools_cache = result
                        return result
                    
                    # 将工具列表转换为按服务器分组的格式
                    tools_by_server = {}
                    for tool in tools_list:
                        tool_name = tool.get("name", "")
                        if ":" in tool_name:
                            server, tool_name_clean = tool_name.split(":", 1)
                        else:
                            server = "default"
                            tool_name_clean = tool_name
                        
                        if server not in tools_by_server:
                            tools_by_server[server] = []
                        
                        tools_by_server[server].append({
                            "name": tool_name_clean,
                            "description": tool.get("description", ""),
                            "inputSchema": tool.get("inputSchema", {})
                        })
                    
                    # 更新缓存
                    MCPConnectionManager._tools_cache = tools_by_server
                    MCPConnectionManager._last_check_time = current_time
                    return tools_by_server
                else:
                    text = await resp.text()
                    if attempt < max_retries - 1:
                        await asyncio.sleep(2 ** attempt)
                        continue
                    result = {"error": f"获取工具失败，HTTP {resp.status}: {text}"}
                    MCPConnectionManager._tools_cache = result
                    return result
        except asyncio.TimeoutError:
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)
                continue
            result = {"error": "连接MCP客户端网关超时"}
            MCPConnectionManager._tools_cache = result
            return result
        except Exception as e:
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)
                continue
            result = {"error": f"连接MCP客户端网关失败: {str(e)}"}
            MCPConnectionManager._tools_cache = result
            return result
    
    result = {"error": "获取工具列表失败，已达到最大重试次数"}
    MCPConnectionManager._tools_cache = result
    return result

def build_react_prompt(tools: dict) -> str:
    """ReactIAgent 系统提示词 - 增强推理版"""
    # 检查是否有错误
    if "error" in tools:
        return f"无法获取工具列表: {tools['error']}\\n\\n请检查MCP客户端网关是否正常运行。"
    
    tool_descs = {}
    for server, items in tools.items():
        if isinstance(items, list):
            for t in items:
                name = t.get("name", "")
                desc = t.get("description", "") or ""
                schema = t.get("inputSchema", {}) or {}
                props = schema.get("properties", {}) if isinstance(schema, dict) else {}
                required = schema.get("required", []) if isinstance(schema, dict) else []
                param_lines = []
                for p, info in props.items():
                    ptype = info.get("type", "unknown") if isinstance(info, dict) else "unknown"
                    pdesc = info.get("description", "") if isinstance(info, dict) else ""
                    param_lines.append(f"- {p} ({ptype}) : {pdesc}")
                param_str = "\\n".join(param_lines) if param_lines else "无参数"
                required_str = ", ".join(required) if required else "无"
                tool_descs[f"{server}:{name}"] = {
                    "描述": desc,
                    "参数": param_str,
                    "必填": required_str,
                }

    sample_json = '{{"tool": "<mcp_server>:<tool_name>", "input": {{...}}}}'

    return (
        "你是 ReactIAgent，负责逐步推理并决定是否调用工具。\\n"
        "你的输出必须严格分为几个部分：\\n"
        "[Think] 详细推理思路（必须包含以下要素）：\\n"
        "  - 分析用户需求的深层含义\\n"
        "  - 识别完成任务所需的关键信息\\n"
        "  - 评估是否需要获取额外信息（如时间、位置等）\\n"
        "  - 选择合适的工具和参数\\n"
        "  - 考虑可能的错误情况和备选方案\\n"
        "[Action] 如果需要工具调用，请只输出 JSON 调用（严格 JSON，无额外解释），格式如下：\\n"
        "```json\\n" + sample_json + "\\n```\\n"
        "如果不需要工具调用，请跳过 [Action]。\\n"
        "[Observation] 当接收到工具返回结果时，将其放在这里。\\n\\n"
        "\\n=== 核心推理原则 ===\\n"
        "1. 信息完整性检查：在调用任何工具前，先检查是否具备所有必要信息\\n"
        "2. 时间敏感性：涉及日期/时间的查询，必须优先调用时间工具获取当前准确时间\\n"
        "3. 工具链思维：考虑是否需要多个工具协作完成任务\\n"
        "4. 错误预判：提前考虑可能出现的错误并准备应对方案\\n"
        "5. 渐进式推理：从简单到复杂，逐步获取和验证信息\\n"
        "\\n=== 强制推理步骤 ===\\n"
        "对于时间相关查询（如'明天'、'下周'等）：\\n"
        "1. 必须首先调用时间工具获取当前准确时间\\n"
        "2. 根据当前时间计算用户请求的具体日期\\n"
        "3. 使用计算后的具体日期调用目标工具\\n"
        "\\n=== 具体推理示例 ===\\n"
        "示例1（时间相关查询）：\\n"
        "用户：查询明天的车票\\n"
        "推理：用户请求涉及'明天'这个相对时间概念，我必须先获取当前准确时间来确定明天的具体日期。\\n"
        "步骤：\\n"
        "1. 调用时间工具获取当前日期\\n"
        "2. 根据当前日期计算明天的具体日期\\n"
        "3. 使用计算后的日期查询车票信息\\n"
        "\\n示例2（位置相关查询）：\\n"
        "用户：查询附近的餐厅\\n"
        "推理：需要先获取当前位置信息，然后查询附近的餐厅\\n"
        "\\n示例3（复杂查询）：\\n"
        "用户：查询从A地到B地的最佳路线\\n"
        "推理：可能需要获取交通信息、天气状况、实时路况等多个信息源\\n"
        "\\n=== 工具使用原则 ===\\n"
        "1. 信息获取优先：在不确定时，优先调用信息获取类工具\\n"
        "2. 参数验证：调用工具前确保参数完整且合理\\n"
        "3. 错误恢复：工具调用失败时，分析原因并尝试替代方案\\n"
        "4. 效率优化：避免不必要的重复调用\\n"
        "\\n重要规则：\\n"
        "1. 不要输出 [FinalAnswer] 或 [Summary]，这些由 SummaryAgent 专门处理\\n"
        "2. 如果任务已完成且无需进一步工具调用，请输出 [任务完成] 标记\\n"
        "3. [Think] 部分必须详细且逻辑清晰，体现完整的推理过程\\n"
        "4. 对于时间相关查询，必须首先调用时间工具，这是强制要求！\\n"
        "\\n可用工具定义：\\n"
        + json.dumps(tool_descs, ensure_ascii=False, indent=2)
        + "\\n\\n请严格遵循以上推理策略和格式要求。"
    )


# ---------- 工具调用 ----------
async def call_mcp_tool(full_tool_name: str, arguments: dict, max_retries: int = 3) -> dict:
    """通过MCP客户端网关调用工具，优化性能和连接池使用"""
    
    for attempt in range(max_retries):
        try:
            # 使用连接管理器检查连接状态
            if not await MCPConnectionManager.check_connection():
                if attempt < max_retries - 1:
                    await asyncio.sleep(2 ** attempt)
                    continue
                return {
                    "success": False,
                    "error": "MCP客户端网关连接失败，请检查服务是否启动"
                }
            
            # 使用连接池获取会话
            session = await MCPConnectionManager._get_session()
            payload = {
                "tool_name": full_tool_name,
                "params": arguments
            }
            
            async with session.post(f"{MCP_CLIENT_URL}/call_tool", json=payload, timeout=60) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return {
                        "success": True,
                        "data": {
                            "content": [{
                                "type": "text",
                                "text": str(result.get("result", ""))
                            }]
                        }
                    }
                else:
                    text = await resp.text()
                    if attempt < max_retries - 1:
                        await asyncio.sleep(2 ** attempt)
                        continue
                    return {
                        "success": False,
                        "error": f"工具调用失败: HTTP {resp.status} - {text}"
                    }
        except asyncio.TimeoutError:
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)
                continue
            return {
                "success": False,
                "error": f"工具调用超时: {full_tool_name}"
            }
        except Exception as e:
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)
                continue
            return {
                "success": False,
                "error": f"连接MCP客户端网关失败: {str(e)}"
            }
    
    return {
        "success": False,
        "error": f"工具调用失败，已达到最大重试次数: {full_tool_name}"
    }


# ---------- 节点 ----------
async def react_agent(state: State):
    """React智能体节点，使用辅助函数减少重复代码"""
    msgs = state["messages"]
    tools = await fetch_tools()
    sys_prompt = build_react_prompt(tools)
    msgs = [{"role": "system", "content": sys_prompt}] + msgs
    response = llm.invoke(msgs)
    return create_success_response(state, response.content)


def summary_agent(state: State):
    """总结智能体节点，使用辅助函数减少重复代码"""
    msgs = state["messages"]
    
    # 获取上下文文件
    context_files = state.get("files", [])
    
    # 构建专业提示词
    sys_prompt = build_summary_prompt(
        state["output_type"],
        get_last_message_content(state) if state["messages"] else "Unknown query",
        context_files
    )
    
    msgs = [{"role": "system", "content": sys_prompt}] + msgs
    response = llm.invoke(msgs)
    
    # 处理输出类型
    final_content, file_path = output_processor.process_output(
        response.content, 
        state["output_type"],
        get_last_message_content(state) if state["messages"] else "Unknown query"
    )
    
    return {
        "messages": [response], 
        "output_type": state["output_type"], 
        "output_file": file_path
    }


# 提取核心内容
def extract_tool_content(result: dict) -> str:
    """从工具响应中提取核心文本内容，提供详细的错误信息"""
    
    # 处理成功响应
    if result.get("success"):
        if 'data' in result and 'content' in result['data']:
            for item in result['data']['content']:
                if item.get('type') == 'text' and 'text' in item:
                    text_content = item['text']
                    
                    # 检查内容是否为空
                    if not text_content or text_content.strip() == "":
                        return "工具调用成功，但返回内容为空"
                    
                    # 尝试解析嵌套的JSON内容
                    try:
                        # 检查是否是JSON字符串
                        if text_content.strip().startswith('{'):
                            nested_data = json.loads(text_content)
                            # 提取实际的文本内容
                            if 'structured_content' in nested_data:
                                content = nested_data['structured_content'].get('result', text_content)
                                return f"工具调用成功: {content}"
                            elif 'data' in nested_data:
                                return f"工具调用成功: {nested_data['data']}"
                            else:
                                return f"工具调用成功: {text_content}"
                        else:
                            return f"工具调用成功: {text_content}"
                    except json.JSONDecodeError:
                        # 如果不是JSON，直接返回文本内容
                        return f"工具调用成功: {text_content}"
                    except Exception as e:
                        return f"工具调用成功但内容解析失败: {str(e)} - 原始内容: {text_content}"
        
        # 检查是否有其他格式的响应
        if 'result' in result:
            return f"工具调用成功: {result['result']}"
        
        return "工具调用成功，但返回内容为空或格式异常"
    
    # 处理错误响应 - 提供更详细的错误信息
    error_msg = result.get("error", "未知错误")
    
    # 分类错误类型
    if "连接失败" in error_msg or "连接MCP" in error_msg:
        return f"❌ 连接错误: {error_msg} (请检查MCP客户端网关服务是否启动)"
    elif "超时" in error_msg:
        return f"⏰ 超时错误: {error_msg} (请检查网络连接或增加超时时间)"
    elif "HTTP" in error_msg:
        return f"🌐 HTTP错误: {error_msg} (请检查MCP服务器状态)"
    elif "重试次数" in error_msg:
        return f"🔄 重试失败: {error_msg} (请稍后重试或检查服务状态)"
    else:
        return f"⚠️ 工具调用失败: {error_msg}"

async def call_tool(state: State):
    """工具调用节点，使用辅助函数减少重复代码"""
    raw_content = get_last_message_content(state)
    
    # 检查是否有JSON格式的工具调用
    json_text = find_json_in_text(raw_content)
    if json_text is None:
        # 如果没有找到JSON，直接跳转到总结节点
        return {
            "messages": [("user", f"[跳转SummaryAgent] {raw_content}")], 
            "output_type": state["output_type"],
            "output_file": state["output_file"],
            "__goto__": "summary"
        }
    
    try:
        # 解析JSON调用
        payload = json.loads(json_text)
        tool_name = payload.get("tool", "")
        args = payload.get("input", {})
        
        if not tool_name:
            return create_error_response(state, "[Observation] ❌ 工具调用失败: 未指定工具名称")
        
        # 检查工具名格式并解析服务器信息
        if ":" not in tool_name:
            # 尝试从可用工具中查找匹配的工具
            tools = await fetch_tools()
            if "error" in tools:
                return create_error_response(state, f"[Observation] ❌ 无法获取工具列表: {tools['error']}")
            
            # 在所有服务器中查找匹配的工具
            found_tool = None
            for server, server_tools in tools.items():
                for tool_info in server_tools:
                    if tool_info.get("name") == tool_name:
                        found_tool = (server, tool_name)
                        break
                if found_tool:
                    break
            
            if not found_tool:
                # 提供可用的工具列表作为调试信息
                available_tools = []
                for server, server_tools in tools.items():
                    for tool_info in server_tools:
                        available_tools.append(f"{server}:{tool_info.get('name')}")
                
                return create_error_response(state, f"[Observation] ❌ 工具不存在: {tool_name}\n可用工具: {', '.join(available_tools[:10])}")
            
            server, tool = found_tool
        else:
            server, tool = tool_name.split(":", 1)
        
        # 构建完整的工具名称
        full_tool_name = f"{server}_{tool}" if server != "default" else tool
        
        # 调用工具
        result = await call_mcp_tool(full_tool_name, args)
        content = extract_tool_content(result)
        
        return create_success_response(state, f"[Observation] {content}")
        
    except json.JSONDecodeError as e:
        return create_error_response(state, f"[Observation] ❌ JSON解析失败: {e}\n原始内容: {json_text}")
    except KeyError as e:
        return create_error_response(state, f"[Observation] ❌ 缺少必要字段: {e}\n请检查工具调用格式")
    except Exception as e:
        return create_error_response(state, f"[Observation] ❌ 工具调用异常: {type(e).__name__}: {str(e)}")


# ---------- 辅助函数 ----------
def get_last_message_content(state: State) -> str:
    """获取最后一条消息内容，减少重复的状态访问"""
    return state["messages"][-1].content.strip()

def create_error_response(state: State, error_message: str) -> dict:
    """创建标准化的错误响应"""
    return {
        "messages": [("user", error_message)],
        "output_type": state["output_type"],
        "output_file": state["output_file"]
    }

def create_success_response(state: State, content: str) -> dict:
    """创建标准化的成功响应"""
    return {
        "messages": [("user", content)],
        "output_type": state["output_type"],
        "output_file": state["output_file"]
    }

# ---------- 构图 ----------
graph_builder = StateGraph(State)
graph_builder.add_node("ReactIAgent", react_agent)
graph_builder.add_node("SummaryAgent", summary_agent)
graph_builder.add_node("tool", call_tool)

graph_builder.add_edge(START, "ReactIAgent")

def route_from_react(state: State):
    """路由函数，使用辅助函数减少重复代码"""
    last = get_last_message_content(state)
    
    # 改进的标记检测逻辑
    if any(marker in last for marker in ["[FinalAnswer]", "[任务完成]"]):
        return "SummaryAgent"
    
    # 检查是否有JSON格式的工具调用
    json_text = find_json_in_text(last)
    if json_text:
        try:
            payload = json.loads(json_text)
            if "tool" in payload:
                return "tool"
        except json.JSONDecodeError as e:
            # 记录JSON解析错误但不中断流程
            print(f"[DEBUG] JSON解析错误: {e}")
        except Exception as e:
            # 记录其他异常但不中断流程
            print(f"[DEBUG] 路由检查异常: {e}")
    
    # 检查是否有明确的结束标记
    if any(marker in last for marker in ["任务完成", "思考结束", "没有更多信息", "无法继续"]):
        return "SummaryAgent"
    
    return "SummaryAgent"

graph_builder.add_conditional_edges("ReactIAgent", route_from_react)
graph_builder.add_edge("tool", "ReactIAgent")
graph_builder.add_edge("SummaryAgent", END)

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)


# ---------- 流式运行 ----------
CONFIG = {"configurable": {"thread_id": "user_001"}}

async def stream_graph_updates(user_input: str, output_type: str = "normal", files: list = []):
    """流式运行图并处理输出，提供美观的用户体验"""
    print(f"\n{'🎯'*3} 多智能体对话开始 {'🎯'*3}")
    print(f"📝 用户输入: {user_input}")
    print(f"📊 输出模式: {output_type}")
    print(f"{'─'*50}")
    
    current_mode = "normal"
    step_count = 0

    # 初始化状态
    initial_state = {
        "messages": [("user", user_input)],
        "output_type": output_type,
        "output_file": None,
        "files": files  # 传递文件数据
    }

    try:
        async for event in graph.astream_events(initial_state, CONFIG, version="v2"):
            step_count += 1
            
            if event["event"] == "on_chat_model_stream":
                token = event["data"]["chunk"].content or ""
                
                # 优化标签检测和输出格式 - 只显示主要标签
                for lbl, emoji, color in [
                    ("[Think]", "💭", "\033[94m"),      # 蓝色思考
                    ("[Action]", "🔧", "\033[93m"),     # 黄色动作
                    ("[Observation]", "📋", "\033[92m"), # 绿色观察
                    ("[FinalAnswer]", "🎯", "\033[92m")  # 绿色最终答案
                ]:
                    if token.startswith(lbl):
                        # 简洁的标签显示，不显示步骤数
                        print(f"\n{color}{emoji} {lbl}\033[0m")
                        token = token[len(lbl):]
                        current_mode = lbl
                        
                        # 如果是FinalAnswer，添加分隔线
                        if lbl == "[FinalAnswer]":
                            print(f"{'─'*50}")
                
                # 优化换行处理
                if token:
                    # 对长文本进行自动换行
                    if current_mode in ["[Observation]", "[FinalAnswer]"]:
                        wrapped_text = textwrap.fill(token, width=80, subsequent_indent='    ')
                        print(f"    {wrapped_text}")
                    else:
                        print(token, end="", flush=True)

            # -------- 工具调用完成 → Observation --------
            if event["event"] == "on_chain_end" and event["name"] == "tool":
                obs = event["data"]["output"]["messages"][-1][1]
                print(f"\n{'📋'} 工具调用结果")
                print(f"    {textwrap.fill(obs, width=80, subsequent_indent='    ')}")

        print(f"\n{'✅'} 任务执行完成")
        print(f"{'─'*50}")
        print("🎉 对话结束！")
        
    except Exception as e:
        print(f"\n{'❌'*3} 任务执行失败 {'❌'*3}")
        print(f"💥 错误类型: {type(e).__name__}")
        print(f"📝 错误信息: {str(e)}")
        print(f"📊 统计信息: 已执行 {step_count} 步推理")
        print(f"{'─'*50}")
        print("⚠️ 对话异常结束")
        raise

# ---------- ReactAgent 类 ----------
class ReactAgent:
    """ReactAgent封装类 - 提供友好的API接口"""
    
    def __init__(self):
        """初始化ReactAgent实例"""
        self.graph = graph
        self.config = CONFIG
        self.conversation_history = []
    
    async def process_query(self, user_input: str, output_type: str = "normal", files: list = None) -> str:
        """
        处理用户查询并返回结果
        
        Args:
            user_input: 用户输入的问题
            output_type: 输出类型 ("normal" 或 "markdown")
            files: 文件列表（可选）
            
        Returns:
            str: 智能体的响应结果
        """
        if files is None:
            files = []
        
        # 记录对话历史
        self.conversation_history.append({
            "timestamp": datetime.now().isoformat(),
            "user_input": user_input,
            "output_type": output_type
        })
        
        # 使用改进的流式更新函数
        print(f"\n{'🚀'*3} ReactAgent 开始处理查询 {'🚀'*3}")
        print(f"📝 用户输入: {user_input}")
        print(f"🎯 输出模式: {output_type}")
        print(f"📊 对话历史: 第 {len(self.conversation_history)} 次对话")
        print(f"{'─'*50}")
        
        # 调用改进的流式处理函数
        result = await stream_graph_updates(user_input, output_type)
        
        print(f"\n{'✅'*3} ReactAgent 处理完成 {'✅'*3}")
        print(f"📊 本次对话已记录到历史")
        
        return result
    
    def get_conversation_stats(self) -> dict:
        """获取对话统计信息"""
        return {
            "total_conversations": len(self.conversation_history),
            "last_conversation": self.conversation_history[-1] if self.conversation_history else None,
            "output_types_used": list(set([conv["output_type"] for conv in self.conversation_history]))
        }
    
    def clear_history(self):
        """清空对话历史"""
        self.conversation_history.clear()
        print("🗑️ 对话历史已清空")
    
    def print_welcome_message(self):
        """打印欢迎信息"""
        print(f"\n{'🌟'*3} ReactAgent 多智能体系统 {'🌟'*3}")
        print("🤖 基于LangGraph的智能协作平台")
        print("📚 支持工具调用、推理思考、报告生成")
        print(f"📊 已处理对话: {len(self.conversation_history)} 次")
        print(f"{'─'*50}")


# ---------- CLI ----------
async def main():
    """主CLI函数，提供友好的用户界面"""
    print(f"\n{'🌟'*3} 多智能体对话系统 {'🌟'*3}")
    print("🤖 基于LangGraph的多智能体协作系统")
    print("📚 支持工具调用、推理思考和专业报告生成")
    print(f"{'─'*50}")
    print("📋 可用命令:")
    print("  • /mode [normal|markdown] - 切换输出模式")
    print("  • /help - 显示帮助信息")
    print("  • /status - 显示当前状态")
    print("  • quit/exit/q - 退出系统")
    print(f"{'─'*50}")
    print("💡 提示: 直接输入问题开始对话")
    
    current_output_type = "normal"
    conversation_count = 0
    
    while True:
        try:
            user_input = input(f"\n{'💬' if conversation_count == 0 else '🔄'} 请输入您的问题: ")
            
            # 处理退出命令
            if user_input.lower() in {"quit", "exit", "q"}:
                print(f"\n{'👋'*3} 感谢使用多智能体对话系统 {'👋'*3}")
                print(f"📊 本次会话共完成 {conversation_count} 次对话")
                print("再见！")
                break
            
            # 处理帮助命令
            if user_input.lower() == "/help":
                print(f"\n{'📖'*3} 帮助信息 {'📖'*3}")
                print("📝 输出模式说明:")
                print("  • normal - 普通文本模式")
                print("  • markdown - 专业Markdown报告模式")
                print("\n🔧 工具调用:")
                print("  • 系统会自动检测并调用合适的工具")
                print("  • 支持文件处理、网络查询等多种工具")
                print("\n💡 使用技巧:")
                print("  • 问题越具体，结果越准确")
                print("  • 可以使用/mode markdown获取专业报告")
                continue
            
            # 处理状态命令
            if user_input.lower() == "/status":
                print(f"\n{'📊'*3} 系统状态 {'📊'*3}")
                print(f"🎯 当前输出模式: {current_output_type}")
                print(f"📈 对话次数: {conversation_count}")
                print(f"🔗 MCP连接状态: 已连接")
                continue
                
            # 检查是否有模式指令
            if user_input.startswith("/mode "):
                parts = user_input.split(" ", 2)
                if len(parts) >= 2:
                    mode = parts[1].lower()
                    if mode in ["normal", "markdown"]:
                        current_output_type = mode
                        print(f"\n🔄 已切换到 {mode} 模式")
                        if len(parts) > 2:
                            user_input = parts[2]
                        else:
                            continue
                    else:
                        print("\n❌ 无效的模式，请使用: normal, markdown")
                        continue
            
            # 处理空输入
            if not user_input.strip():
                print("\n⚠️ 请输入有效的问题")
                continue
            
            # 执行对话
            conversation_count += 1
            print(f"\n{'🚀'*3} 开始第 {conversation_count} 次对话 {'🚀'*3}")
            await stream_graph_updates(user_input, current_output_type)
            
        except (EOFError, KeyboardInterrupt):
            print(f"\n{'👋'*3} 感谢使用多智能体对话系统 {'👋'*3}")
            print(f"📊 本次会话共完成 {conversation_count} 次对话")
            print("再见！")
            break


async def run_main():
    """运行主程序并确保资源正确清理"""
    try:
        await main()
    finally:
        # 确保清理所有异步资源
        await MCPConnectionManager.cleanup()

if __name__ == "__main__":
    asyncio.run(run_main())